/*
 * Copyright (c) 2012, Pasi Luostarinen
 * All rights reserved.
 * 
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met: 
 * 
 * 1. Redistributions of source code must retain the above copyright notice, this
 *    list of conditions and the following disclaimer. 
 * 2. Redistributions in binary form must reproduce the above copyright notice,
 *    this list of conditions and the following disclaimer in the documentation
 *    and/or other materials provided with the distribution. 
 * 
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
 * ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
 * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 * 
 * The views and conclusions contained in the software and documentation are those
 * of the authors and should not be interpreted as representing official policies, 
 * either expressed or implied, of the FreeBSD Project.
 */
package fi.serv.btk.bencode;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.UnsupportedEncodingException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;

public class BEncode
{

  /**
   * Is thrown when parsing failure has occured
   */
  public static class ParsingFailed extends Exception
  {

    public ParsingFailed(String pMsg)
    {
      super(pMsg);
    }
  }

  /**
   * Is thrown when given object cannot be converted to BEncode
   */
  public static class InvalidObject extends Exception
  {

    public InvalidObject(String pMsg)
    {
      super(pMsg);
    }
  }

  /**
   * An item of a decoded bencode string
   */
  public static class BItem
  {

    private final Object mObj;
    private String mStrCache = null;

    private BItem(Object pObj)
    {
      mObj = pObj;
    }

    /**
     *
     * @return the class of stored item
     */
    public Class getClassOfItem()
    {
      return (mObj.getClass());
    }

    /**
     *
     * @return The object without casting
     */
    public Object getObject()
    {
      return (mObj);
    }

    /**
     * Gives the string object if object can be converted to string
     *
     * @return the string object
     * @throws dimo.util.BEncoder.BEncode.ParsingFailed Is thrown when object
     * cannot be converted to a string
     */
    public String getString() throws ParsingFailed
    {
      if(mStrCache == null)
      {
        try
        {
          if(mObj instanceof byte[])
          {
            mStrCache = new String((byte[]) mObj, "UTF-8");
          }
          else
          {
            throw new ParsingFailed("Cannot convert to string. The object is " + mObj.getClass().getName());
          }
        }
        catch(UnsupportedEncodingException ex)
        {
          assert (false);
        }
      }

      return mStrCache;
    }

    /**
     * Gives the raw bytes if object can be converted to it
     *
     * @return the byte[] object
     * @throws dimo.util.BEncoder.BEncode.ParsingFailed Is thrown when object
     * cannot be converted to a byte array
     */
    public byte[] getBytes() throws ParsingFailed
    {
      if(mObj instanceof byte[])
      {
        return (byte[]) mObj;
      }
      else
      {
        throw new ParsingFailed("Cannot convert to byte[]. The object is " + mObj.getClass().getName());
      }
    }

    /**
     * Gives the integer object if object can be converted to integer
     *
     * @return the integer object
     * @throws dimo.util.BEncoder.BEncode.ParsingFailed Is thrown when object
     * cannot be converted to a integer
     */
    public Integer getInteger() throws ParsingFailed
    {
      if(mObj instanceof Long)
      {
        return ((Long) mObj).intValue();
      }
      else
      {
        throw new ParsingFailed("Cannot convert to Interger. The object is " + mObj.getClass().getName());
      }
    }

    /**
     * Gives the long object if object can be converted to long
     *
     * @return the long object
     * @throws dimo.util.BEncoder.BEncode.ParsingFailed Is thrown when object
     * cannot be converted to a long
     */
    public Long getLong() throws ParsingFailed
    {
      if(mObj instanceof Long)
      {
        return (Long) mObj;
      }
      else
      {
        throw new ParsingFailed("Cannot convert to Long. The object is " + mObj.getClass().getName());
      }
    }

    /**
     * Gives the list object if object can be converted to list
     *
     * @return the list object
     * @throws dimo.util.BEncoder.BEncode.ParsingFailed Is thrown when object
     * cannot be converted to a list
     */
    public List<BItem> getList() throws ParsingFailed
    {
      if(mObj instanceof List)
      {
        @SuppressWarnings("unchecked")
        List<BItem> out = (List<BItem>) mObj;
        return (out);
      }
      else
      {
        throw new ParsingFailed("Cannot convert to List. The object is " + mObj.getClass().getName());
      }
    }

    /**
     * Gives the map object if object can be converted to map
     *
     * @return the map object
     * @throws dimo.util.BEncoder.BEncode.ParsingFailed Is thrown when object
     * cannot be converted to a map
     */
    public Map<String, BItem> getDict() throws ParsingFailed
    {
      if(mObj instanceof Map)
      {
        @SuppressWarnings("unchecked")
        Map<String, BItem> out = (Map<String, BItem>) mObj;
        return (out);
      }
      else
      {
        throw new ParsingFailed("Cannot convert to Dictionary. The object is " + mObj.getClass().getName());
      }
    }

    /**
     * if object is list then this method gives the pInxed item of a list
     *
     * @param pIndex the index of list
     * @return the object
     * @throws dimo.util.BEncoder.BEncode.ParsingFailed is thrown when object
     * cannot be converted to a list
     */
    public BItem get(int pIndex) throws ParsingFailed
    {
      return (getList().get(pIndex));
    }

    /**
     * if object is list then this returns the size of the list
     *
     * @return the size of the list
     * @throws dimo.util.BEncoder.BEncode.ParsingFailed is thrown when object
     * cannot be converted to a list
     */
    public int getListSize() throws ParsingFailed
    {
      return (getList().size());
    }

    /**
     * get a object from a map if the object is map
     *
     * @param pKey the key of the item
     * @return the object
     * @throws dimo.util.BEncoder.BEncode.ParsingFailed is thrown when object
     * cannot be converted to a Map
     */
    public BItem get(String pKey) throws ParsingFailed
    {
      return (getDict().get(pKey));
    }
  }

  /**
   * A encoded BEncode array
   */
  static public class EncodedItem
  {

    private final byte[] mData;

    private EncodedItem(byte[] pData)
    {
      mData = pData;
    }

    public byte[] getBytes()
    {
      return (mData);
    }
  }

  /**
   * encodes object is a BEncode array
   *
   * @param pStream the stream where the output is written
   * @param pObj the object to be converted
   * @throws IOException is thrown when there is a IO problem
   * @throws dimo.util.BEncoder.BEncode.InvalidObject is thrown when encoder
   * encountered object that cannot be converted to BEncode
   */
  public static void encode(OutputStream pStream, Object pObj) throws IOException, InvalidObject
  {
    selectEncoder(pStream, pObj);
  }

  /**
   * encodes object is a BEncode array
   *
   * @param pObj the object to be converted
   * @return result BEncode string
   * @throws IOException is thrown when there is a IO problem
   * @throws dimo.util.BEncoder.BEncode.InvalidObject is thrown when encoder
   * encountered object that cannot be converted to BEncode
   */
  public static EncodedItem encode(Object pObj) throws IOException, InvalidObject
  {
    ByteArrayOutputStream tmp = new ByteArrayOutputStream();
    encode(tmp, pObj);
    tmp.flush();

    return (new EncodedItem(tmp.toByteArray()));
  }

  private static void encodeInt(OutputStream pStream, List<?> pObj) throws InvalidObject
  {
    try
    {
      pStream.write('l');

      for(Object i : pObj)
      {
        selectEncoder(pStream, i);
      }

      pStream.write('e');
    }
    catch(IOException ex)
    {
      assert (false);
    }
  }

  private static void encodeInt(OutputStream pStream, String pObj) throws IOException
  {
    try
    {
      byte[] tmp = pObj.getBytes("UTF-8");
      byte[] len = (Integer.toString(tmp.length, 10) + ":").getBytes("UTF-8");
      pStream.write(len);
      pStream.write(tmp);
    }
    catch(UnsupportedEncodingException ex)
    {
      assert (false);
    }

  }

  private static void encodeInt(OutputStream pStream, byte[] pObj) throws IOException
  {
    try
    {
      byte[] len = (Integer.toString(pObj.length, 10) + ":").getBytes("UTF-8");
      pStream.write(len);
      pStream.write(pObj);
    }
    catch(UnsupportedEncodingException ex)
    {
      assert (false);
    }
  }

  private static void encodeInt(OutputStream pStream, Map<String, ?> pObj) throws IOException, InvalidObject
  {
    try
    {

      ArrayList<byte[]> keys = new ArrayList<byte[]>();
      for(String i : pObj.keySet())
      {
        keys.add(i.getBytes("UTF-8"));
      }

      Collections.sort(keys, new Comparator<byte[]>()
      {
        @Override
        public int compare(byte[] o1, byte[] o2)
        {
          int len = Math.min(o1.length, o2.length);
          for(int i = 0; i < len; ++i)
          {
            if(o1[i] == o2[i])
            {
              continue;
            }
            else if(o1[i] < o2[i])
            {
              return -1;
            }
            else
            {
              return 1;
            }
          }

          if(o1.length == o2.length)
          {
            return 0;
          }
          else if(o1.length < o2.length)
          {
            return -1;
          }
          else
          {
            return 1;
          }
        }
      });

      pStream.write('d');

      for(byte[] i : keys)
      {
        Object obj = pObj.get(new String(i, "UTF-8"));
        assert obj != null;

        encodeInt(pStream, i);
        selectEncoder(pStream, obj);
      }

      pStream.write('e');
    }
    catch(UnsupportedEncodingException ex)
    {
      assert false;
    }
  }

  private static void encodeInt(OutputStream pStream, int pObj) throws IOException
  {
    try
    {
      String tmp = "i" + Integer.toString(pObj, 10) + "e";
      pStream.write(tmp.getBytes("UTF-8"));
    }
    catch(UnsupportedEncodingException ex)
    {
      assert (false);
    }
  }

  private static void encodeInt(OutputStream pStream, long pObj) throws IOException
  {
    try
    {
      String tmp = "i" + Long.toString(pObj, 10) + "e";
      pStream.write(tmp.getBytes("UTF-8"));
    }
    catch(UnsupportedEncodingException ex)
    {
      assert (false);
    }
  }

  private static void encodeInt(OutputStream pStream, EncodedItem pObj) throws IOException
  {
    pStream.write(pObj.getBytes());
  }

  @SuppressWarnings("unchecked")
  private static void selectEncoder(OutputStream pStream, Object pObj) throws IOException, InvalidObject
  {
    if(pObj instanceof Integer)
    {
      encodeInt(pStream, (Integer) pObj);
    }
    else if(pObj instanceof Long)
    {
      encodeInt(pStream, (Long) pObj);
    }
    else if(pObj instanceof String)
    {
      encodeInt(pStream, (String) pObj);
    }
    else if(pObj instanceof byte[])
    {
      encodeInt(pStream, (byte[]) pObj);
    }
    else if(pObj instanceof List)
    {
      encodeInt(pStream, (List) pObj);
    }
    else if(pObj instanceof Map)
    {
      encodeInt(pStream, (Map) pObj);
    }
    else if(pObj instanceof EncodedItem)
    {
      encodeInt(pStream, (EncodedItem) pObj);
    }
    else
    {
      throw new InvalidObject("Invalid object was encountered, type was: " + pObj.getClass().toString());
    }
  }

  /**
   * decodes the stream
   * @param pStream the stream
   * @return BEncode item which contains all the data
   * @throws IOException is thrown when read error has encountered
   * @throws fi.serv.btk.bencode.BEncode.ParsingFailed is thrown when parsing has failed
   */
  public static BItem decode(InputStream pStream) throws IOException, ParsingFailed
  {

    char character = getChar(pStream);
    return decode(pStream, character);

  }

  private static BItem decode(InputStream pStream, char pChar) throws IOException, ParsingFailed
  {
    if(pChar == 'i')
    {
      return decodeInteger(pStream);
    }
    else if(pChar == 'l')
    {
      return decodeList(pStream);
    }
    else if(pChar == 'd')
    {
      return decodeDict(pStream);
    }
    else
    {
      return decodeString(pStream, pChar);
    }
  }

  private static byte[] getEndOfBlock(InputStream pStream) throws IOException, ParsingFailed
  {
    ByteArrayOutputStream out = new ByteArrayOutputStream();

    while(true)
    {
      char data = getChar(pStream);

      if(data == 'e')
      {
        break;
      }

      out.write(data);
    }
    return out.toByteArray();
  }

  private static BItem decodeInteger(InputStream pStream) throws IOException, ParsingFailed
  {
    byte[] buffer = getEndOfBlock(pStream);

    return new BItem(Long.parseLong(new String(buffer, "UTF-8"), 10));
  }

  private static BItem decodeList(InputStream pStream) throws IOException, ParsingFailed
  {
    List<BItem> out = new ArrayList<BItem>();

    while(true)
    {
      char chr = getChar(pStream);
      if(chr == 'e')
      {
        break;
      }

      out.add(decode(pStream, chr));
    }

    return (new BItem(out));
  }

  private static BItem decodeDict(InputStream pStream) throws IOException, ParsingFailed
  {
    Map<String, BItem> out = new TreeMap<String, BItem>();


    while(true)
    {
      char chr = getChar(pStream);
      if(chr == 'e')
      {
        break;
      }

      String key = decodeString(pStream, chr).getString();
      BItem value = decode(pStream);
      out.put(key, value);
    }

    return new BItem(out);
  }

  private static char getChar(InputStream pStream) throws IOException, ParsingFailed
  {
    int val = pStream.read();
    if(val < 0)
    {
      throw new ParsingFailed("Unexpected end of file");
    }

    return ((char) val);
  }

  private static BItem decodeString(InputStream pStream, char pFirstChar) throws ParsingFailed, IOException
  {
    if('0' > pFirstChar || pFirstChar > '9')
    {
      throw new ParsingFailed("Unexpected character encountered");
    }

    ByteArrayOutputStream bstream = new ByteArrayOutputStream();
    bstream.write(pFirstChar);

    while(true)
    {
      char chr = getChar(pStream);

      if(chr == ':')
      {
        break;
      }

      if('0' > chr || chr > '9')
      {
        throw new ParsingFailed("Unexpected character encountered");
      }

      bstream.write(chr);
    }

    int len = Integer.parseInt(new String(bstream.toByteArray(), "UTF-8"), 10);

    byte[] buffer = new byte[len];
    if(pStream.read(buffer) != len)
    {
      throw new ParsingFailed("Unexpected end of file");
    }

    return new BItem(buffer);
  }
}
